# 加载cifar10数据集
from torchvision import transforms
from torch.utils.data import DataLoader, Dataset
import os
from PIL import Image
import numpy as np
import glob

label_name = ["airplane",
              "automobile",
              "bird",
              "cat",
              "deer",
              "dog",
              "frog",
              "horse",
              "ship",
              "truck"]

label_dict = {}

for idx, name in enumerate(label_name):
    label_dict[name] = idx


# 是一个用于加载图片的函数，它使用 PIL 库来读取图像，并将图像转换为 RGB 格式。
def default_loader(path):
    return Image.open(path).convert("RGB")

# 对图片的预处理
# 训练数据的预处理操作，包括随机裁剪、水平翻转、垂直翻转、随机旋转、随机灰度化和颜色抖动
train_transform = transforms.Compose([
    transforms.RandomResizedCrop((28, 28)),
    transforms.RandomHorizontalFlip(),
    transforms.RandomVerticalFlip(),
    transforms.RandomRotation(90),
    transforms.RandomGrayscale(0.1),
    transforms.ColorJitter(0.3, 0.3, 0.3, 0.3),
    transforms.ToTensor()
])


class MyDataset(Dataset):
    '''
        im_list : 图像路径列表
        transform : 图像预处理
        loader ： 图像加载函数
    '''
    def __init__(self, im_list,
                 transform=None,
                 loader=default_loader):
        super(MyDataset, self).__init__()
        imgs = []

        for im_item in im_list:
            print(im_item)
            im_label_name = im_item.split("\\")[-2]
            imgs.append([im_item, label_dict[im_label_name]])

            self.imgs = imgs
            self.transform = transform
            self.loader = loader

    def __getitem__(self, index):
        '''
         根据索引 index 获取对应位置的图像路径和标签。
        :param index:
        :return:
        '''
        im_path, im_label = self.imgs[index]
        im_data = self.loader(im_path)
        if self.transform is not None:
            im_data = self.transform(im_data)

        return im_data, im_label

    def __len__(self):
        return len(self.imgs)


im_train_list = glob.glob("./train/*/*.png")
im_test_list = glob.glob("./test/*/*.png")

train_dataset = MyDataset(im_train_list, transform=train_transform)
test_dataset = MyDataset(im_test_list, transform=train_transform)

train_loader = DataLoader(dataset=train_dataset,
                               batch_size=6,
                               shuffle=True,
                               num_workers=4)
test_loader = DataLoader(dataset=test_dataset,
                              batch_size=6,
                              shuffle=False,
                              num_workers=4)

print("num_of_train", len(train_dataset))
print("num_of_test", len(test_dataset))
